import torch
import torch.nn as nn

from utils import idx2onehot


class VAE(nn.Module):

    def __init__(self, encoder_layer_sizes, latent_size, decoder_layer_sizes,
                 conditional=False, num_labels=0): #这是一个条件变分自编码器

        super().__init__() #初始化父类

        if conditional: #如果是条件变分自编码器
            assert num_labels > 0 #确保有标签

        assert type(encoder_layer_sizes) == list  #确保encoder_layer_sizes是一个列表
        assert type(latent_size) == int  #隐藏层的大小 是int
        assert type(decoder_layer_sizes) == list

        self.latent_size = latent_size

        self.encoder = Encoder(   #初始化编码器
            encoder_layer_sizes, latent_size, conditional, num_labels)
        self.decoder = Decoder(   #初始化解码器
            decoder_layer_sizes, latent_size, conditional, num_labels)

    def forward(self, x, c=None):

        if x.dim() > 2:  #如果输入的维度大于2，则将其展平为28*28的向量
            x = x.view(-1, 28*28)

        means, log_var = self.encoder(x, c)  #通过编码器得到均值和方差
        z = self.reparameterize(means, log_var)  #通过重参数化得到隐变量
        recon_x = self.decoder(z, c)  #通过解码器得到重构的图像
        return recon_x, means, log_var, z  #返回重构的图像，均值，方差，隐变量

    def reparameterize(self, mu, log_var):  #重参数化

        std = torch.exp(0.5 * log_var)  #计算方差
        eps = torch.randn_like(std)  #生成噪声
        return mu + eps * std  #返回重参数化的隐变量

    def inference(self, z, c=None):
        recon_x = self.decoder(z, c)  #通过解码器得到重构的图像
        return recon_x  #返回重构的图像


class Encoder(nn.Module):
    def __init__(self, layer_sizes, latent_size, conditional, num_labels):

        super().__init__()  #初始化父类

        self.conditional = conditional  #是否是条件变分自编码器
        if self.conditional:
            layer_sizes[0] += num_labels  #如果是条件变分自编码器，则将标签的维度加到输入层
        self.MLP = nn.Sequential()  #初始化一个空的 Sequential 容器，
        # 并将其赋值给 self.MLP 属性。这个容器可以按顺序添加多个网络层，从而构建一个多层感知器（MLP）。

        for i, (in_size, out_size) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):  #遍历输入输出的大小
            self.MLP.add_module(   #添加一个线性层和一个激活层
                name="L{:d}".format(i), module=nn.Linear(in_size, out_size))  #线性层的输入和输出大小分别为 in_size 和 out_size
            self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())  #激活层使用 ReLU 激活函数

        self.linear_means = nn.Linear(layer_sizes[-1], latent_size)  #线性层的输入和输出大小分别为 layer_sizes[-1] 和 latent_size
        self.linear_log_var = nn.Linear(layer_sizes[-1], latent_size)  #线性层的输入和输出大小分别为 layer_sizes[-1] 和 latent_size

    def forward(self, x, c=None):

        if self.conditional:  #如果是条件变分自编码器
            c = idx2onehot(c, n=10)  #将标签转换为 one-hot 编码
            x = torch.cat((x, c), dim=-1)  #将输入和标签拼接在一起
        x = self.MLP(x)  #通过 MLP 得到隐变量
        means = self.linear_means(x)  #通过线性层得到均值
        log_vars = self.linear_log_var(x)  #通过线性层得到方差
        return means, log_vars  #返回均值和方差


class Decoder(nn.Module):

    def __init__(self, layer_sizes, latent_size, conditional, num_labels):

        super().__init__()

        self.MLP = nn.Sequential()
        self.conditional = conditional
        if self.conditional:
            input_size = latent_size + num_labels
        else:
            input_size = latent_size
        for i, (in_size, out_size) in enumerate(zip([input_size]+layer_sizes[:-1], layer_sizes)):
            self.MLP.add_module(   #添加一个线性层和一个激活层
                name="L{:d}".format(i), module=nn.Linear(in_size, out_size))  #线性层的输入和输出大小分别为 in_size 和 out_size
            if i+1 < len(layer_sizes):  #如果不是最后一层
                self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())  #激活层使用 ReLU 激活函数
            else:
                self.MLP.add_module(name="sigmoid", module=nn.Sigmoid())  #最后一层使用 sigmoid 激活函数

    def forward(self, z, c):

        if self.conditional:
            c = idx2onehot(c, n=10)  #将标签转换为 one-hot 编码
            z = torch.cat((z, c), dim=-1)  #将隐变量和标签拼接在一起
        x = self.MLP(z)
        return x
